/*
 * Copyright © 2020-2021 尛飛俠（Denvie） All rights reserved.
 */

package cn.denvie.springboot.mybatis.generator;

import org.mybatis.generator.api.IntrospectedColumn;
import org.mybatis.generator.api.IntrospectedTable;
import org.mybatis.generator.api.dom.java.*;
import org.mybatis.generator.config.PropertyRegistry;
import org.mybatis.generator.internal.DefaultCommentGenerator;
import org.mybatis.generator.internal.util.StringUtility;

import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.List;
import java.util.Optional;
import java.util.Properties;

/**
 * 自定义注释生成器。
 *
 * @author Denvie
 * @date 2021/7/3
 * @since 1.0.0
 */
public class MyCommentGenerator extends DefaultCommentGenerator {
    private SimpleDateFormat dateFormat = null;
    private String author;
    private String version;
    private String copyright;

    /**
     * 设置用户配置的参数
     */
    @Override
    public void addConfigurationProperties(Properties properties) {
        super.addConfigurationProperties(properties);
        this.dateFormat = new SimpleDateFormat(properties.getProperty(
                PropertyRegistry.COMMENT_GENERATOR_DATE_FORMAT, "yyyy-MM-dd"));
        this.author = properties.getProperty("author", System.getProperty("user.name"));
        this.version = properties.getProperty("version", "1.0.0");
        this.copyright = properties.getProperty("copyright",
                "Copyright © 2020-2021 尛飛俠（Denvie） All rights reserved.");
    }

    /**
     * 添加Java文件头注释
     */
    @Override
    public void addJavaFileComment(CompilationUnit compilationUnit) {
        // Java文件头加入版权信息
        GeneratorUtils.addFileComment(compilationUnit, copyright);
    }

    /**
     * 添加Model类注释
     */
    @Override
    public void addModelClassComment(TopLevelClass topLevelClass, IntrospectedTable introspectedTable) {
        // 为Model声明 import 信息
        if (!topLevelClass.getType().getFullyQualifiedName().endsWith("Example")
                && !topLevelClass.getType().getFullyQualifiedName().endsWith("Mapper")
                && !topLevelClass.getType().getFullyQualifiedName().endsWith("Dao")) {
            List<IntrospectedColumn> allColumns = introspectedTable.getAllColumns();
            boolean hasIdentityOrAutoIncrement = false;
            for (IntrospectedColumn column : allColumns) {
                if (column.isIdentity() || column.isAutoIncrement()) {
                    hasIdentityOrAutoIncrement = true;
                    break;
                }
            }
            if (hasIdentityOrAutoIncrement) {
                topLevelClass.addImportedType(new FullyQualifiedJavaType(
                        "com.baomidou.mybatisplus.annotation.IdType"));
            }
            topLevelClass.addImportedType(new FullyQualifiedJavaType(
                    "com.baomidou.mybatisplus.annotation.TableField"));
            if (hasIdentityOrAutoIncrement) {
                topLevelClass.addImportedType(new FullyQualifiedJavaType(
                        "com.baomidou.mybatisplus.annotation.TableId"));
            }
            topLevelClass.addImportedType(new FullyQualifiedJavaType(
                    "com.baomidou.mybatisplus.annotation.TableName"));
            topLevelClass.addImportedType(new FullyQualifiedJavaType(
                    "io.swagger.annotations.ApiModel"));
            topLevelClass.addImportedType(new FullyQualifiedJavaType(
                    "io.swagger.annotations.ApiModelProperty"));
            topLevelClass.addImportedType(new FullyQualifiedJavaType(
                    "lombok.Data"));
            topLevelClass.addImportedType(new FullyQualifiedJavaType(
                    "lombok.EqualsAndHashCode"));
            topLevelClass.addImportedType(new FullyQualifiedJavaType(
                    "lombok.NoArgsConstructor"));
        }

        // 给Model实体类添加注释
        String remarks = introspectedTable.getRemarks();
        if (!StringUtility.stringHasValue(remarks)) {
            remarks = introspectedTable.getFullyQualifiedTable().getIntrospectedTableName();
        } else {
            remarks = remarks + "对象";
        }
        addComments(topLevelClass, remarks, true);
        // 给Model实体类添加Lombok注解
        topLevelClass.addAnnotation("@Data");
        topLevelClass.addAnnotation("@EqualsAndHashCode(callSuper = true)");
        topLevelClass.addAnnotation("@NoArgsConstructor");
        // 给Model实体类添加 @Table 注解
        topLevelClass.addAnnotation("@TableName(\""
                + introspectedTable.getFullyQualifiedTableNameAtRuntime() + "\")");
        // 添加@ApiModel注解
        topLevelClass.addAnnotation("@ApiModel(\"" + remarks + "\")");
    }

    /**
     * 添加字段注释
     */
    @Override
    public void addFieldComment(Field field, IntrospectedTable introspectedTable,
                                IntrospectedColumn introspectedColumn) {
        String remarks = introspectedColumn.getRemarks();
        if (!StringUtility.stringHasValue(remarks)) {
            remarks = introspectedColumn.getActualColumnName();
        }
        // 给字段添加注释
        addComments(field, remarks, false);
        // 给主键字段添加 @TableId 注解
        if (introspectedColumn.isIdentity()) {
            field.addAnnotation("@TableId(value = \"" + introspectedColumn.getActualColumnName()
                    + "\", type = IdType.INPUT)");
        }
        // 给自增字段添加 @TableId 注解
        else if (introspectedColumn.isAutoIncrement()) {
            field.addAnnotation("@TableId(value = \"" + introspectedColumn.getActualColumnName()
                    + "\", type = IdType.AUTO)");
        }
        // 给普通字段添加 @Column 注解
        else {
            field.addAnnotation("@TableField(value = \"" + introspectedColumn.getActualColumnName() + "\")");
        }
        // 添加@ApiModelProperty注解
        field.addAnnotation("@ApiModelProperty(\"" + remarks + "\")");
    }

    /**
     * add class or field comment.
     */
    private void addComments(JavaElement element, String remarks, boolean isClass) {
        // 数据库中特殊字符需要转义
        if (remarks.contains("\"")) {
            remarks = remarks.replace("\"", "'");
        }
        element.addJavaDocLine("/**");
        // 获取数据库字段的备注信息
        String[] remarkLines = remarks.split(System.getProperty("line.separator"));
        for (String remarkLine : remarkLines) {
            element.addJavaDocLine(" * " + remarkLine);
        }
        if (isClass) {
            element.addJavaDocLine(" * ");
            element.addJavaDocLine(" * @author " + author);
            element.addJavaDocLine(" * @date " + dateFormat.format(new Date()));
            element.addJavaDocLine(" * @since " + version);
        }
        element.addJavaDocLine(" */");
    }

    /**
     * 添加普通方法的注释，主要是XxxMapper接口方法的注释
     */
    @Override
    public void addGeneralMethodComment(Method method, IntrospectedTable introspectedTable) {
        method.addJavaDocLine("/**");
        method.addJavaDocLine(" * " + method.getName());
        // 添加 @param 注释
        List<Parameter> parameters = method.getParameters();
        if (parameters != null && parameters.size() > 0) {
            method.addJavaDocLine(" * ");
            for (Parameter parameter : parameters) {
                method.addJavaDocLine(" * @param " + parameter.getName() + " " + parameter.getType().getShortName());
            }
        }
        // 添加 @return 注释
        Optional<FullyQualifiedJavaType> returnTypeOptional = method.getReturnType();
        returnTypeOptional.ifPresent(fullyQualifiedJavaType ->
                method.addJavaDocLine(" * @return " + fullyQualifiedJavaType.getShortName()));
        // 添加 @throws 注释
        List<FullyQualifiedJavaType> exceptions = method.getExceptions();
        if (exceptions != null && exceptions.size() > 0) {
            for (FullyQualifiedJavaType exception : exceptions) {
                method.addJavaDocLine(" * @throws " + exception.getShortName() + " " + exception.getShortName());
            }
        }
        method.addJavaDocLine(" */");
    }
}
